import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

from langchain.schema import (
    SystemMessage,
    HumanMessage,
    AIMessage
)


# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL")
model = os.getenv("DEFAULT_MODEL")
max_tokens = os.getenv("MAX_TOKENS_DEFAULT")

llm = ChatOpenAI(
    model=model,
    temperature=0,
    #max_tokens=None,  这个不能加，加了会报错
    #timeout=None,
    max_retries=2,
    api_key=api_key,
    base_url=base_url
)

# messages = [
#     (
#         "system",
#         "You are a helpful assistant that translates English to French. Translate the user sentence.",
#     ),
#     (
#         "human",
#         "I love programming."
#     ),
# ]
#
# ai_msg = llm.invoke(messages)
# print(ai_msg.content)

# 翻译
# translate_prompt = ChatPromptTemplate.from_messages(
#     [
#         (
#             "system",
#             "You are a helpful assistant that translates {input_language} to {output_language}.",
#         ),
#         ("human", "{input}"),
#     ]
# )
#
# chain = translate_prompt | llm
# translate_res = chain.invoke(
#     {
#         "input_language": "English",
#         "output_language": "Chinese",
#         "input": "I love my family",
#     }
# )
# print(translate_res.content)

# 起名字
# name_prompt = ChatPromptTemplate.from_template("你是一个起名大师,请模仿示例起3个{country}名字,比如男孩经常被叫做{boy},女孩经常被叫做{girl}")
# name_prompt = ChatPromptTemplate.from_template(role = "起名大师", template="请模仿示例起3个{country}名字,比如男孩经常被叫做{boy},女孩经常被叫做{girl}")
# message = {
#     "country" : "日本特色的",
#     "boy" : "狗蛋",
#     "girl" : "翠花"
# }
# chain = name_prompt | llm
# name_res = chain.invoke(message)
# print(name_res.content)

# 用角色Message的方式调用
messages = [
    SystemMessage(content="You are a helpful assistant."),
    HumanMessage(content="Knock knock."),
    AIMessage(content="Who's there."),
    HumanMessage(content="Orange"),
]

res = llm(messages)
print("模型回复：", res.content)